package org.ns.tomdog.server.support;

import com.sun.xml.internal.messaging.saaj.util.ByteOutputStream;
import org.ns.tomdog.server.ServerResponse;

import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.Charset;
import java.util.*;

public class ServerResponseImpl implements ServerResponse {
    private int status = 200;
    private String version = "HTTP/1.1";
    private String message = "OK";
    private Map<String, List<String>> headers = new HashMap<>();
    private OutputStream outputStream;
    private ByteOutputStream cacheout = new ByteOutputStream(1024);
    private boolean headout = false;

    public ServerResponseImpl(OutputStream outputStream) {
        this.outputStream = outputStream;
    }

    @Override
    public void setVersion(String version) {
        this.version = version;
    }

    @Override
    public void setStatus(int status) {
        this.status = status;
    }

    @Override
    public void setError(String error) {
        this.message = error;
    }

    @Override
    public void setHeader(String name, String value) {
        this.headers.put(name, Arrays.asList(value));
    }

    @Override
    public void addHeader(String name, String value) {
        List<String> values = this.headers.get(name);
        if (values == null) {
            values = new ArrayList<>();
            this.headers.put(name, values);
        }
        values.add(value);
    }

    @Override
    public void setContentType(String contentType) {
        setHeader("contentType", contentType);
    }

    @Override
    public void serContentLength(String contentLength) {
        setHeader("contentLength", contentLength);
    }

    private byte[] readyWriteBytes() {
        byte[] bytes = null;
        if (!headout) {
            this.headout = true;
            StringBuilder stringBuilder = new StringBuilder();
            stringBuilder.append(version).append(" ").append(status).append(" ").append(message).append("\r\n");
            // headers
            for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
                List<String> values = entry.getValue();
                for (String value : values) {
                    stringBuilder.append(entry.getKey()).append(":").append(value).append("\r\n");
                }
            }
            stringBuilder.append("\r\n");
            byte[] headBytes = stringBuilder.toString().getBytes(Charset.forName("UTF-8"));
            int count = cacheout.getCount();
            bytes = new byte[headBytes.length + count];
            System.arraycopy(headBytes, 0, bytes, 0, headBytes.length);
            System.arraycopy(cacheout.getBytes(), 0, bytes, headBytes.length, count);
        } else {
            bytes = cacheout.toByteArray();
        }

        cacheout.reset();
        return bytes;
    }


    @Override
    public void write(byte[] bytes) {
        cacheout.write(bytes);
    }

    @Override
    public void write(byte[] bytes, int off, int len) {
        cacheout.write(bytes, off, len);
    }

    @Override
    public void flush() throws IOException {
        byte[] bytes = readyWriteBytes();
        if (bytes != null && bytes.length > 0) {
            outputStream.write(bytes);
        }
        outputStream.flush();
    }

    @Override
    public void close() throws IOException {
        flush();
        outputStream.close();
    }
}
